// Copyright 2012 The golab Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package machine

import (
	"fmt"
	"log"
	"sync"

	"../cmd"
)

// A Group defines an interior node of the tree structure representing the
// network of machines to be used for an experiment.  A group may have both
// Subgroups and Members where members are Client leaves.  If a Group is
// recreated from a serialized format (JSON) one should almost always fix the
// structure by calling FixIDs.
//
// Ideally group IDs should be globally unique but this is not strictly
// necessary.  
type Group struct {
	ID        string             // An ID for this group 
	SubGroups map[string]*Group  // map of Groups where the key should be the Group's ID
	Members   map[string]*Client // map of Clients where the key should be the Client's Host
	mux       sync.RWMutex       // a mutex to protect for the SubGroups and Members map
}

// NewGroup creates a new Group node with a given ID.
func NewGroup(id string) *Group {
	g := &Group{
		ID:        id,
		SubGroups: make(map[string]*Group),
		Members:   make(map[string]*Client),
	}
	return g
}

// flatten reduces a multi level group by moving all members of subgroups to the
// top level group.  It also removes any duplicate members at the same time.
func (g *Group) flatten() []*Client {
	fg := &Group{
		ID:      g.ID,
		Members: make(map[string]*Client),
	}
	g.flattenHelp(fg)
	rg := make([]*Client, len(fg.Members))
	i := 0
	for _, v := range fg.Members {
		rg[i] = v
		i++
	}
	return rg
}

// flattenHelp is a recursive friendly helper function for flatten.
func (g *Group) flattenHelp(fg *Group) {
	g.mux.RLock()
	for _, v := range g.SubGroups {
		v.flattenHelp(fg)
	}
	g.mux.RUnlock()
	fg.mux.Lock()
	for k, v := range g.Members {
		fg.Members[k] = v
	}
	fg.mux.Unlock()
}

// FixIDs is used to assign Group IDs to the key in the parent group's map.
func (g *Group) FixIDs(l *log.Logger, p string) {
	g.mux.Lock()
	defer g.mux.Unlock()
	for k, v := range g.SubGroups {
		v.ID = k
		l.Printf("ok: FixIDs: (GroupID=%s SubGroupID=%s)", g.ID, v.ID)
	}
	for k, v := range g.Members {
		if v.Host == "" {
			v.Host = k
			v.ResPath = fmt.Sprintf("%s/%s", p, v.Host)
		}
		l.Printf("ok: FixIDs: (GroupID=%s MemberID=%s)", g.ID, v.Host)
	}
	for _, v := range g.SubGroups {
		v.FixIDs(l, p)
	}
}

// GroupAdd attaches a given subgroup.
func (g *Group) GroupAdd(sg *Group) {
	g.mux.Lock()
	defer g.mux.Unlock()
	g.SubGroups[sg.ID] = sg
}

// GroupRm removes a subgroup with the given id.
func (g *Group) GroupRm(id string) {
	g.mux.Lock()
	defer g.mux.Unlock()
	delete(g.SubGroups, id)
}

// ClientAdd attaches a given member Client.
func (g *Group) ClientAdd(c *Client) {
	g.mux.Lock()
	defer g.mux.Unlock()
	g.Members[c.Host] = c
}

// ClientRm removes a member Client with the given id.
func (g *Group) ClientRm(host string) {
	g.mux.Lock()
	defer g.mux.Unlock()
	delete(g.Members, host)
}

// Find returns a pointer to the Group with a matching ID.
func (g *Group) Find(ID string) *Group {
	if g.ID == ID {
		return g
	}
	g.mux.RLock()
	defer g.mux.RUnlock()
	for k, v := range g.SubGroups {
		if ID == k {
			return v
		}
	}
	for _, v := range g.SubGroups {
		r := v.Find(ID)
		if r != nil {
			return r
		}
	}
	return nil
}

// Start will begin an asynchronous command on a Group.
func (g *Group) Start(c *cmd.Cmd) error {
	g.mux.RLock()
	defer g.mux.RUnlock()
	for _, v := range g.SubGroups {
		err := v.Start(c)
		if err != nil {
			return err
		}
	}
	for _, v := range g.Members {
		err := v.Start(c)
		if err != nil {
			return err
		}
	}
	return nil
}

// Start will begin an asynchronous command on a Group with a given target.
//
// See Client.StartTarget for the details of how targets are used.
func (g *Group) StartTarget(c *cmd.Cmd, tg *Group) error {
	cl := g.flatten()
	tl := tg.flatten()
	if len(cl) > len(tl) {
		return fmt.Errorf("Target group too small (len=%d)", len(tl))
	}
	for i := range cl {
		err := cl[i].StartTarget(c, tl[i].Host)
		if err != nil {
			return err
		}
	}
	return nil
}

// Wait will call Client.Wait(cmdID) on all the member Clients.
func (g *Group) Wait(cmdID string) error {
	g.mux.RLock()
	defer g.mux.RUnlock()
	for _, v := range g.SubGroups {
		err := v.Wait(cmdID)
		if err != nil {
			return err
		}
	}
	for _, v := range g.Members {
		_, err := v.Wait(cmdID)
		if err != nil {
			return err
		}
	}
	return nil
}

// Stop will call Client.Stop(cmdID) on all the member Clients.
func (g *Group) Stop(cmdID string) error {
	g.mux.RLock()
	defer g.mux.RUnlock()
	for _, v := range g.SubGroups {
		err := v.Stop(cmdID)
		if err != nil {
			return err
		}
	}
	for _, v := range g.Members {
		_, err := v.Stop(cmdID)
		if err != nil {
			return err
		}
	}
	return nil
}

// Kill will call Client.Kill(cmdID) on all the member Clients.
func (g *Group) Kill(cmdID string) error {
	g.mux.RLock()
	defer g.mux.RUnlock()
	for _, v := range g.SubGroups {
		err := v.Kill(cmdID)
		if err != nil {
			return err
		}
	}
	for _, v := range g.Members {
		_, err := v.Kill(cmdID)
		if err != nil {
			return err
		}
	}
	return nil
}

// Save will call Client.Save(cmdID, stepID) on all the member Clients.
func (g *Group) Save(cmdID, stepID string) error {
	g.mux.RLock()
	defer g.mux.RUnlock()
	for _, v := range g.SubGroups {
		err := v.Save(cmdID, stepID)
		if err != nil {
			return err
		}
	}
	for _, v := range g.Members {
		err := v.Save(cmdID, stepID)
		if err != nil {
			return err
		}
	}
	return nil
}

// SaveAll will call Client.SaveAll() on all the member Clients.
func (g *Group) SaveAll() error {
	g.mux.RLock()
	defer g.mux.RUnlock()
	for _, v := range g.SubGroups {
		err := v.SaveAll()
		if err != nil {
			return err
		}
	}
	for _, v := range g.Members {
		err := v.SaveAll()
		if err != nil {
			return err
		}
	}
	return nil
}

// Close will call Client.Close() on all the member Clients.  It also logs the
// closing of the clients in the experiment level logger.
func (g *Group) Close(l *log.Logger) {
	g.mux.Lock()
	defer g.mux.Unlock()
	for _, v := range g.Members {
		err := v.Close()
		if err != nil {
			l.Printf("ERROR: Close:v.Close: %v (GroupID=%s MemberID=%s)", err, g.ID, v.Host)
			continue
		}
		l.Printf("ok: Close:v.Close: (GroupID=%s MemberID=%s)", g.ID, v.Host)
	}
	for _, v := range g.SubGroups {
		v.Close(l)
	}
	l.Printf("ok: Close: (GroupID=%s)", g.ID)
}
